import json
import os
import time

import javalang
from app.logic.utils.filepath import AST_DIR, SRC_DIR

'''
按项目路径生成抽象语法树
'''


def extract_ast(java_file_path):
    """
    提取单个 Java 文件的 AST
    """
    try:
        with open(java_file_path, 'r', encoding='utf-8') as f:
            code = f.read()
        tree = javalang.parse.parse(code)
        return tree
    except Exception as e:
        print(f"Failed to parse {java_file_path}: {e}")
        return None


class GenerateAST:
    """
    初始化
    work_dir: 工作目录
    project_path: 项目路径
    """

    def __init__(self, work_path_dir, project_path_dir):
        self.workspace = os.path.join(work_path_dir, str(project_path_dir))
        self.src_dir = os.path.join(work_path_dir, str(project_path_dir), SRC_DIR)
        self.ast_dir = os.path.join(work_path_dir, str(project_path_dir), AST_DIR)

        # 确保 ast 存储目录存在
        if not os.path.exists(self.ast_dir):
            os.makedirs(self.ast_dir)

    def serialize_ast(self, node):
        """
        将 AST 节点序列化为字典格式（递归）
        """
        if isinstance(node, list):
            return [self.serialize_ast(child) for child in node]
        elif isinstance(node, javalang.ast.Node):
            return {
                'type': type(node).__name__,
                'children': self.serialize_ast(list(node.children))
            }
        else:
            return str(node)

    def generate_ast_for_file(self, java_file_path):
        """
        生成单个文件的 AST 并保存为 JSON 文件
        """
        ast = extract_ast(java_file_path)
        if ast:
            serialized_ast = self.serialize_ast(ast)
            # 保存 AST 到 JSON 文件
            # relative_path = os.path.relpath(java_file_path, self.src_dir)
            # ast_file = os.path.join(self.ast_dir, f"{relative_path}.json")
            # 使用文件名而非相对路径保存 AST 文件
            file_name = os.path.basename(java_file_path)  # 获取文件名
            ast_file = os.path.join(self.ast_dir, f"{file_name}.json")
            os.makedirs(os.path.dirname(ast_file), exist_ok=True)
            with open(ast_file, 'w', encoding='utf-8') as f:
                json.dump(serialized_ast, f, indent=4, ensure_ascii=False)
            print(f"AST saved to {ast_file}")

    def generate_ast_for_directory(self):
        """
        遍历目录中的所有 Java 文件并生成它们的 AST
        """
        for dir_name in os.listdir(self.src_dir):
            project_dir_path = os.path.join(self.src_dir, dir_name)
            if not os.path.isdir(project_dir_path):
                continue

            # 遍历项目中的 Java 文件
            for root, _, files in os.walk(project_dir_path):
                for file in files:
                    if file.endswith('.java'):
                        java_file = os.path.join(root, file)
                        self.generate_ast_for_file(java_file)

    def extract_and_save_ast(self):
        """
        打印开始时间，生成 AST 并保存，最后打印结束时间和总耗时
        """
        start_time = time.time()
        print("AST Extraction Start Time: " + time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()))

        # 生成 AST
        self.generate_ast_for_directory()

        # 打印结束时间
        print("AST extraction completed at " + time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()))

        # 计算总耗时
        total_time = time.time() - start_time
        print("Total time: " + str(total_time // 3600) + "h" + str((total_time % 3600) // 60) + "m" + str(
            (total_time % 3600) % 60) + "s")


def gen_ast(work_dir, project_path):
    ast_generator = GenerateAST(work_dir, project_path)
    ast_generator.extract_and_save_ast()


# 使用方法示例
if __name__ == "__main__":
    work_dir = '../../../data'  # 请修改为实际路径
    project_path = '29'  # 请修改为实际项目路径

    ast_generator = GenerateAST(work_dir, project_path)
    ast_generator.extract_and_save_ast()
